'''
Audio Recognition Task Client
这个服务用于接收语音识别任务，并将识别结果发送给调用方。
可以运行于没有公网 IP 的环境，通过 WebSocket 连接到任务分发服务器。

terrence@tenclass.com
2024-09
'''
import pyaudio
import wave
import opuslib
import os
import sys
import json
import uuid
import time
import queue
import asyncio
import logging
import threading
import websocket
import base64
import hashlib
from typing import Dict, Any
from collections import deque
from util import is_segment
from util import get_string_no_punctuation_or_emoji,remove_punctuation_and_length
from concurrent.futures import ThreadPoolExecutor, TimeoutError

sys.stdout.reconfigure(encoding='utf-8')  # ✅ 让 print() 正确处理 UTF-8 字符

import numpy as np
import tts
from logger import setup_logging
from settings import load_config

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', stream=sys.stdout)
# 加载环境变量
from dotenv import load_dotenv
load_dotenv()
# 常量定义
TTS_TASK_SERVER_URL = "ws://localhost:8084"
#TTS_TASK_SERVER_URL = os.getenv("TTS_TASK_SERVER_URL")


class TtsWorker:
    def __init__(self, wsapp, session_id, config, audio_param):
        self.wsapp = wsapp
        self.session_id = session_id
        self.config = config
        self.audio_param = audio_param
        # 客户端状态相关
        self.client_abort = False
        # tts相关变量
        self.tts_first_text = None
        self.tts_last_text = None
        self.tts_start_speak_time = None
        self.tts_duration = 0
        self.speaking = False

        # 线程任务相关
        self.loop = None
        self.stop_event = threading.Event()
        self.tts_queue = queue.Queue()
        self.executor = ThreadPoolExecutor(max_workers=5)
        self.scheduled_tasks = deque()
        logging.info(f'TTS agent : {self.config["selected_module"]["TTS"]}')
        self.tts = tts.create_instance(
                self.config["selected_module"]["TTS"],
                self.config["TTS"][self.config["selected_module"]["TTS"]],
                self.config["delete_audio"],
                self.audio_param
            )
        threading.Thread(target=self.start_event_loop, daemon=True).start()
        self.tts_priority = threading.Thread(target=self._priority_thread, daemon=True)
        self.tts_priority.start()

    def start_event_loop(self):
        try:
            self.loop = asyncio.get_running_loop()
        except RuntimeError:
            print("Create new_event_loop")
            self.loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self.loop)
            self.loop.run_forever() 

    def _priority_thread(self):
        while not self.stop_event.is_set():
            text = None
            try:
                future = self.tts_queue.get()
                text = None
                try:
                    tts_file,cnt, text, state = future.result(timeout=10)
                    if os.path.exists(tts_file):
                        opus_datas, duration = self.tts.wav_to_opus_data(tts_file)
                    else:
                        opus_datas = []
                        duration = 0
                except TimeoutError:
                    logging.error("TTS 任务超时")
                    continue
                except Exception as e:
                    logging.error(f"TTS 任务出错: {e}")
                    continue
                if not self.client_abort:
                    # 如果没有中途打断就发送语音
                    print(f'发送 opus_datas list :{type(opus_datas)}  length :{len(opus_datas)}')
                    asyncio.run_coroutine_threadsafe(
                        self.sendAudioMessage(opus_datas, duration, cnt, text, state), self.loop
                    )
                    time.sleep(duration-0.2)
                if self.tts.delete_audio_file and os.path.exists(tts_file):
                    os.remove(tts_file)

            except Exception as e:
                self.clearSpeakStatus()
                asyncio.run_coroutine_threadsafe(
                    self.wsapp.send(json.dumps({"type": "tts", "timeStamp": int(time.time() * 1000), "state": "stop", "session_id": self.session_id})),
                    self.loop
                )
                logging.error(f"tts_priority priority_thread: {text}{e}")

    def remove_unpronounced_symbols(self, text):
        # 定义需要删除的符号
        symbols_to_remove = ['*', '#']
        for symbol in symbols_to_remove:
            text = text.replace(symbol, '')
        return text
            
    def speak_and_play(self, cnt, state, text):
        if text is None or len(text) <= 0:
            logging.info(f"无需tts转换，query为空，{text}")
            return None
        tts_file = self.tts.to_tts(self.remove_unpronounced_symbols(text))
        if tts_file is None:
            logging.error(f"tts转换失败，{text}")
            return None
        logging.debug(f"TTS 文件生成完毕")
        return tts_file, cnt, text, state

    def clearSpeakStatus(self):
        logging.debug(f"清除服务端讲话状态")
        self.asr_server_receive = True
        self.tts_last_text = None
        self.tts_first_text = None
        self.tts_duration = 0
        self.tts_start_speak_time = None

    async def sendAudioMessage(self, audios, duration, cnt, text, state):
        base_delay = self.tts_duration
        if state == 'start':
            self.tts_start_speak_time = time.time()
            self.wsapp.send(json.dumps({
                "type": "tts",
                "timeStamp": int(time.time() * 1000),
                "state": "start",
                "sample_rate": self.audio_param['sample_rate'],
                "session_id": self.session_id
            }))
        # # 调度文字显示任务
        self.send_sentence_start(text)
        self.tts_duration += duration
        # 发送音频数据
        for opus_packet in audios:
            if not isinstance(opus_packet, bytes):
                raise ValueError("opus_packet 必须是 bytes 类型")
            self.wsapp.send(json.dumps({
                "type": "audio",
                "timeStamp": int(time.time() * 1000),
                "text": base64.b64encode(opus_packet).decode('utf-8'),
                "session_id": self.session_id
            }))
        self.send_sentence_end(text)
        print(f'发送 {int(time.time() * 1000)} - {duration}/s 音频数据 end')
        if state == 'end' or cnt == -1:
            stop_duration = self.tts_duration - (time.time() - self.tts_start_speak_time)
            time.sleep(duration)
            stop_task = asyncio.create_task(
                self.schedule_with_interrupt(stop_duration, self.send_tts_stop(text))
            )
            self.scheduled_tasks.append(stop_task)

    def send_sentence_start(self, text):
        message = {
            "type": "tts",
            "timeStamp": int(time.time() * 1000),
            "state": "sentence_start",
            "text": text,
            "session_id": self.session_id
        }
        self.wsapp.send(json.dumps(message))
        logging.info(message)

    def send_sentence_end(self, text):
        message = {
            "type": "tts",
            "timeStamp": int(time.time() * 1000),
            "state": "sentence_end",
            "text": text,
            "session_id": self.session_id
        }
        self.wsapp.send(json.dumps(message))
        logging.info(message)

    def send_tts_stop(self, text):
        message_stop = {
            "type": "tts",
            "timeStamp": int(time.time() * 1000),
            "state": "stop",
            "session_id": self.session_id
        }
        self.wsapp.send(json.dumps(message_stop))
        logging.info(message_stop)
        self.clearSpeakStatus()

    async def schedule_with_interrupt(delay, coro):
        """可中断的延迟调度"""
        try:
            await asyncio.sleep(delay)
            await coro
        except asyncio.CancelledError:
            pass
    
    def start(self, mode):
        self.speaking = True
    
    def stop(self):
        self.tts_queue.queue.clear()
        self.speaking = False


class TtsTaskClient:
    def __init__(self):
        self.workers = {}
        self.wsapp = None
        self.audio_param = None
        setup_logging()  # 最先初始化日志
        self.config = load_config()

    def initialize(self):
        logging.info("task client init")

    def get_worker(self, session_id):
        if session_id not in self.workers:
            self.workers[session_id] = TtsWorker(self.wsapp, session_id, self.config, self.audio_param)
        return self.workers[session_id]

    def parse_text_message(self, message):
        data = json.loads(message)
        print("Rec : ", data)
        session_id = data['session_id']
        if data['type'] == 'llm':
            worker = self.get_worker(session_id)
            state = data['state']
            if state == 'start' or state == 'next' or state == 'end':
                logging.info(f"Recv tts word: {state}-{data['content']}")
                future = worker.executor.submit(worker.speak_and_play, data['cnt'], state, data['content'])
                worker.tts_queue.put(future)
            logging.info(f'Worker {session_id} started {state}')
        elif data['type'] == 'hello':
            self.audio_param = data['audio_params']
            print(f'Hello config audio param: {self.audio_param}')
        elif data['type'] == 'finish':
            if session_id in self.workers:
                del self.workers[session_id]
                logging.info(f'Worker {session_id} finished')
        else:
            logging.warning(f'Unknown message type: {data["type"]}')

    def on_message(self, wsapp, message):
        try:
            if isinstance(message, str):
                self.parse_text_message(message)
            else:
                print("Message error not support binary message")
        except Exception as e:
            logging.error(f"An error occurred: {e}", exc_info=True)

    def on_open(self, wsapp):
        logging.info('Connected to TTS Task Server.')

    async def run(self):
        logging.info('Starting TTS Task Client...')
        self.wsapp = websocket.WebSocketApp(
            TTS_TASK_SERVER_URL,
            on_message=self.on_message,
            on_open=self.on_open
        )

        while True:
            try:
                ret = self.wsapp.run_forever()
                if ret is None:
                    break
            except Exception as e:
                logging.error(f"An error occurred: {e}", exc_info=True)
            # Remove all workers
            self.workers = {}
            logging.info('Reconnecting to the TTS Task Server in 3 seconds...')
            time.sleep(3)


if __name__ == "__main__":
    task_client = TtsTaskClient()
    task_client.initialize()
    asyncio.run(task_client.run())

